import numpy as np
import matplotlib.pyplot as plt
class LineProcessor:
    """处理3D直线拟合和异面直线最近点计算的类"""
    
    def __init__(self):
        self.tolerance = 1e-6  # 平行判断的容差
        self.debug = True     # 是否开启调试信息
    
    def get_dict_fitted_points(self, edges_3d):
        """
        对edges_3d中的每个边的点集合进行直线拟合，返回拟合后的点集字典，
        并计算left_edge和bottom_edge（或top_edge）的最近点中点作为"center"点
        
        参数:
            edges_3d: dict, 形如{"left_edge": [pt1, pt2, ...], ...}
        返回:
            dict, 形如{"left_edge": [拟合点1, 拟合点2, ...], ..., "center": center_point}
        """
        fitted_dict = {}
        
        # 拟合每条边的点
        for edge_name, points in edges_3d.items():
            if isinstance(points, (list, tuple)) and len(points) > 0:
                fitted_points = self.fitted_points(points)
                fitted_dict[edge_name] = fitted_points
            else:
                fitted_dict[edge_name] = []

        # 获取left_edge和bottom_edge的拟合点
        left_points = fitted_dict.get("left_edge", [])
        bottom_points = fitted_dict.get("bottom_edge", [])
        top_points = fitted_dict.get("top_edge", [])
        
        pts1 = self._get_valid_points(left_points)
        
        # 优先使用bottom_edge计算center
        if len(bottom_points) > 0:
            pts2 = self._get_valid_points(bottom_points)
            if len(pts1) >= 2 and len(pts2) >= 2:
                center_point = self.skew_lines_shortest_midpoint_by_points(pts1[0], pts1[-1], pts2[0], pts2[-1])
                fitted_dict["center"] = center_point
                return fitted_dict
        
        # 如果bottom_edge不存在或无效，尝试使用top_edge计算center
        if len(top_points) > 0:
            pts2 = self._get_valid_points(top_points)
            if len(pts1) >= 2 and len(pts2) >= 2:
                center_point = self.skew_lines_shortest_midpoint_by_points(pts1[0], pts1[-1], pts2[0], pts2[-1])
                fitted_dict["center"] = center_point
                return fitted_dict
        
        # 如果都无法计算，返回None
        fitted_dict["center"] = None
        return fitted_dict

    
    def get_dict_fitted_vector(self, fitted_dict):
        """
        根据拟合点集字典计算每条边的起点和单位向量，并增加四个角点（如存在）
        
        参数:
            fitted_dict: 包含拟合点集和中心点的字典
            
        返回:
            dict: 每条边对应的起点和单位向量，以及四个角点
        """
        
        # 结果字典
        vector_dict = {}
        
        # 1. 获取中心点
        center_point = fitted_dict.get("center")
        
        # 2. 处理left_edge
        left_start_point = None
        if "left_edge" in fitted_dict and len(fitted_dict["left_edge"]) >= 2:
            left_points = fitted_dict["left_edge"]
            # 在前5个点中找到距离center最近的点作为起点
            start_index = self._find_closest_point_index(left_points[:5], center_point)
            left_start_point = left_points[start_index]
            # 计算候选向量并投票确定方向
            unit_vec = self._determine_direction_by_voting(left_points, start_index)
            vector_dict["left_edge"] = {"start_point": left_start_point, "unit_vector": unit_vec}
        
        # 3. 处理bottom_edge
        bottom_start_point = None
        if "bottom_edge" in fitted_dict and len(fitted_dict["bottom_edge"]) >= 2:
            bottom_points = fitted_dict["bottom_edge"]
            # 在前5个点中找到距离center最近的点作为起点
            start_index = self._find_closest_point_index(bottom_points[:5], center_point)
            bottom_start_point = bottom_points[start_index]
            # 计算候选向量并投票确定方向
            unit_vec = self._determine_direction_by_voting(bottom_points, start_index)
            vector_dict["bottom_edge"] = {"start_point": bottom_start_point, "unit_vector": unit_vec}
        
        # 4. 处理top_edge（新增）
        top_start_point = None
        if "top_edge" in fitted_dict and len(fitted_dict["top_edge"]) >= 2:
            top_points = fitted_dict["top_edge"]
            # 在前5个点中找到距离center最近的点作为起点
            start_index = self._find_closest_point_index(top_points[:5], center_point)
            top_start_point = top_points[start_index]
            # 计算候选向量并投票确定方向
            unit_vec = self._determine_direction_by_voting(top_points, start_index)
            vector_dict["top_edge"] = {"start_point": top_start_point, "unit_vector": unit_vec}
        
        vector_dict["center"] = center_point

        # 增加"left_bottom"为left和bottom起点的中点
        if left_start_point is not None and bottom_start_point is not None:
            left_bottom = (np.array(left_start_point) + np.array(bottom_start_point)) / 2.0
            vector_dict["left_bottom"] = left_bottom
        
        # 增加"left_top"为left和top起点的中点（新增）
        if left_start_point is not None and top_start_point is not None:
            left_top = (np.array(left_start_point) + np.array(top_start_point)) / 2.0
            vector_dict["left_top"] = left_top

        return vector_dict


    def _find_closest_point_index(self, points, target_point):
        """在点集中找到距离目标点最近的点索引"""
        import numpy as np
        min_dist = float('inf')
        min_index = 0
        
        for i, point in enumerate(points):
            if point is None:
                continue
            dist = np.linalg.norm(np.array(point) - np.array(target_point))
            if dist < min_dist:
                min_dist = dist
                min_index = i
                
        return min_index

    def _determine_direction_by_voting(self, points, start_index):
        """通过投票机制确定单位向量方向"""
        import numpy as np
        from collections import Counter
        
        # 候选终点索引
        end_indices = [
            start_index + 1, 
            start_index + 4,
            start_index + 7,
            start_index + 10,
            start_index + 13
        ]
        
        # 有效候选向量
        candidate_vectors = []
        
        for end_idx in end_indices:
            # 确保索引在有效范围内
            if 0 <= end_idx < len(points) and points[end_idx] is not None:
                start_pt = np.array(points[start_index])
                end_pt = np.array(points[end_idx])
                vec = end_pt - start_pt
                
                # 跳过零向量
                norm = np.linalg.norm(vec)
                if norm > 1e-6:
                    unit_vec = vec / norm
                    candidate_vectors.append(unit_vec)
        
        # 如果没有候选向量，返回默认向量
        if not candidate_vectors:
            if start_index < len(points) - 1:
                vec = np.array(points[start_index + 1]) - np.array(points[start_index])
                return vec / np.linalg.norm(vec)
            else:
                return np.array([0, 0, 0])
        
        # 投票确定主要方向
        vote_counter = Counter()
        
        for i, vec_i in enumerate(candidate_vectors):
            for j, vec_j in enumerate(candidate_vectors):
                if i < j:
                    dot_product = np.dot(vec_i, vec_j)
                    # 方向相同
                    if dot_product > 0.9:
                        vote_counter[i] += 1
                    # 方向相反
                    elif dot_product < -0.9:
                        vote_counter[i] -= 1
        
        # 选择票数最高的方向
        if vote_counter:
            main_idx = max(vote_counter, key=vote_counter.get)
            return candidate_vectors[main_idx]
        
        # 默认返回第一个候选向量
        return candidate_vectors[0]
    def _get_valid_points(self, edge):
        """
        获取有效的3D点
        
        参数:
            edge: 点列表
        返回:
            list: 有效的3D点列表
        """
        return [pt for pt in edge if pt is not None and hasattr(pt, '__len__') and len(pt) == 3]
    
    def skew_lines_shortest_midpoint_by_points(self,A1, A2, B1, B2):
        """
        求由两组点表示的异面直线间最短距离线段的中点坐标。
        返回: 中点坐标（长度为3的ndarray）
        """
        A1 = np.array(A1)
        A2 = np.array(A2)
        B1 = np.array(B1)
        B2 = np.array(B2)
        u = A2 - A1
        v = B2 - B1
        w0 = A1 - B1
        a = np.dot(u, u)
        b = np.dot(u, v)
        c = np.dot(v, v)
        d = np.dot(u, w0)
        e = np.dot(v, w0)
        denom = a * c - b * b
        if denom == 0:
            # 平行或重合，返回某一条直线上的点
            return (A1 + B1) / 2
        s = (b * e - c * d) / denom
        t = (a * e - b * d) / denom
        P = A1 + s * u  # 直线1上的点
        Q = B1 + t * v  # 直线2上的点
        midpoint = (P + Q) / 2
        return midpoint
    
    def plot_skew_lines_and_shortest(self,A1, A2, B1, B2):
        """
        可视化两条异面直线及最短距离线段和中点
        """
        A1 = np.array(A1)
        A2 = np.array(A2)
        B1 = np.array(B1)
        B2 = np.array(B2)
        # 计算最短距离线段的端点和中点
        u = A2 - A1
        v = B2 - B1
        w0 = A1 - B1
        a = np.dot(u, u)
        b = np.dot(u, v)
        c = np.dot(v, v)
        d = np.dot(u, w0)
        e = np.dot(v, w0)
        denom = a * c - b * b
        if denom == 0:
            P = A1
            Q = B1
        else:
            s = (b * e - c * d) / denom
            t = (a * e - b * d) / denom
            P = A1 + s * u
            Q = B1 + t * v
        midpoint = (P + Q) / 2
        fig = plt.figure()
        ax = fig.add_subplot(111, projection='3d')
        # 画直线1
        ax.plot([A1[0], A2[0]], [A1[1], A2[1]], [A1[2], A2[2]], 'b-', label='Line 1')
        # 画直线2
        ax.plot([B1[0], B2[0]], [B1[1], B2[1]], [B1[2], B2[2]], 'g-', label='Line 2')
        # 画最短距离线段
        ax.plot([P[0], Q[0]], [P[1], Q[1]], [P[2], Q[2]], 'r--', label='Shortest Segment')
        # 画中点
        ax.scatter(midpoint[0], midpoint[1], midpoint[2], c='r', marker='o', s=80, label='Midpoint')
        ax.legend()
        ax.set_xlabel('X')
        ax.set_ylabel('Y')
        ax.set_zlabel('Z')
        plt.title('Skew Lines and Shortest Segment')
        plt.show()
        return midpoint

    
    def fitted_points(self, points):
        """
        对大致在一条直线上的3D点集进行直线拟合，将所有点投影到拟合直线上
        参数:
            points: list of np.array - 3D点集 [[x1, y1, z1], [x2, y2, z2], ...]
        返回:
            list of np.array - 所有点投影到拟合直线上的点（与原始点数量一致，顺序对应）
        """
        if len(points) < 3:
            # 点太少无法有效拟合
            return points

        # 过滤掉None点，记录原始索引
        valid_idx_and_points = [(i, p) for i, p in enumerate(points) if p is not None]
        if len(valid_idx_and_points) < 3:
            # 有效点太少，直接返回有效点
            return [p for _, p in valid_idx_and_points]

        valid_indices = [i for i, _ in valid_idx_and_points]
        valid_points = [p for _, p in valid_idx_and_points]
        points_arr = np.array(valid_points)

        # 步骤1: 使用RANSAC拟合直线模型
        n_samples = len(points_arr)
        best_inliers = []
        best_line = None

        # RANSAC参数
        max_iterations = 100
        threshold = 0.02  # 2厘米的阈值（根据点云尺度调整）
        min_inliers = max(3, n_samples // 3)  # 最小内点数

        tried = 0
        found = False
        while tried < 5 and not found:
            for _ in range(max_iterations):
                # 随机选择两点确定直线
                idx1, idx2 = np.random.choice(n_samples, 2, replace=False)
                p1, p2 = points_arr[idx1], points_arr[idx2]
                direction = p2 - p1

                # 跳过距离太近的点
                if np.linalg.norm(direction) < 1e-5:
                    continue

                # 计算所有点到直线的距离
                # 距离公式: ||(p - p1) × (p - p2)|| / ||p2 - p1||
                vec1 = points_arr - p1
                vec2 = points_arr - p2
                cross_prod = np.cross(vec1, vec2)
                distances = np.linalg.norm(cross_prod, axis=1) / np.linalg.norm(direction)

                # 统计内点
                inliers = np.where(distances < threshold)[0]

                if len(inliers) > len(best_inliers) and len(inliers) >= min_inliers:
                    best_inliers = inliers
                    best_line = (p1.copy(), p2.copy())
                    found = True
            if not found:
                if threshold * 2 <= 0.1:  # 最大阈值设为0.1
                    threshold *= 2  # 增大阈值再试
                tried += 1
        # 如果没有找到有效模型，返回原始点
        if best_line is None:
            return points

        # 使用所有内点重新估计直线
        inlier_points = points_arr[best_inliers]
        centroid = np.mean(inlier_points, axis=0)

        # 计算主方向 (PCA)
        centered = inlier_points - centroid
        cov_matrix = centered.T @ centered
        eig_vals, eig_vectors = np.linalg.eig(cov_matrix)
        direction = eig_vectors[:, np.argmax(eig_vals)]  # 最大特征值对应的特征向量
        direction = direction / np.linalg.norm(direction)

        # 步骤2: 将所有有效点（非None）投影到拟合直线上
        projected_points = []
        for p in valid_points:
            vec = p - centroid
            scalar = np.dot(vec, direction)
            proj_p = centroid + scalar * direction
            projected_points.append(proj_p)

        # 构造与原始输入顺序一致的输出（None点仍为None，其余为投影点）
        result = []
        proj_idx = 0
        for i in range(len(points)):
            if i in valid_indices:
                result.append(projected_points[proj_idx])
                proj_idx += 1
            else:
                result.append(None)
        return result